#!/usr/bin/env python3
"""
GMAE 诊断脚本
在服务器上运行，验证GMAE修复效果
"""

import torch
import torch.nn as nn
import sys
import os

# 添加项目路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

def diagnose_gmae():
    """诊断GMAE模块"""
    print("=== GMAE 诊断报告 ===")

    try:
        from src.modules import GMAE_Decoder, GMAE

        # 1. 检查解码器
        print("\n1. GMAE解码器检查:")
        decoder = GMAE_Decoder(emb_dim=128, n_features=350)
        print(f"   ✅ 解码器创建成功")
        print(f"   - 参数数量: {sum(p.numel() for p in decoder.parameters())}")
        print(f"   - 可训练参数: {sum(p.numel() for p in decoder.parameters() if p.requires_grad)}")

        # 测试前向传播
        z_u = torch.randn(32, 128)
        z_c = torch.randn(350, 128)
        recon = decoder(z_u, z_c)
        print(f"   ✅ 前向传播成功")
        print(f"   - 输入形状: {z_u.shape}, {z_c.shape}")
        print(f"   - 输出形状: {recon.shape}")
        print(f"   - 输出范围: [{recon.min().item():.4f}, {recon.max().item():.4f}]")

        # 2. 检查梯度
        print("\n2. 梯度检查:")
        loss = recon.mean()
        loss.backward()

        has_gradients = False
        for name, param in decoder.named_parameters():
            if param.grad is not None and param.grad.abs().sum() > 0:
                has_gradients = True
                print(f"   ✅ {name}: 有梯度 (范数: {param.grad.norm().item():.6f})")
                break

        if not has_gradients:
            print("   ❌ 没有检测到梯度!")

        # 3. 检查GMAE完整模块
        print("\n3. GMAE完整模块检查:")

        # 模拟数据
        class MockData:
            def __init__(self):
                self.n_users = 100
                self.n_items = 200
                self.n_entities = 300
                self.graph_ui_edge_index = torch.zeros((2, 10), dtype=torch.long)
                self.graph_ui_edge_weight = torch.ones(10)
                self.graph_ik_edge_index = torch.zeros((2, 10), dtype=torch.long)
                self.graph_ik_edge_weight = torch.ones(10)
                self.graph_ii_edge_index = torch.zeros((2, 10), dtype=torch.long)
                self.graph_ii_edge_weight = torch.ones(10)
                self.graph_uk_edge_index = torch.zeros((2, 10), dtype=torch.long)
                self.graph_uk_edge_weight = torch.ones(10)
                self.graph_ku_edge_index = torch.zeros((2, 10), dtype=torch.long)
                self.graph_ku_edge_weight = torch.ones(10)
                self.user_interest_features = torch.randn(100, 350)

        data = MockData()

        # 创建GMAE
        gmae = GMAE(
            gcn_layers=3,
            mlp_layers=2,
            n_features=350,
            n_nodes=100,
            n_entities=300,
            emb_dim=128,
            gcn_edge_index_uk=data.graph_uk_edge_index,
            gcn_edge_weight_uk=data.graph_uk_edge_weight,
            gcn_edge_index_ku=data.graph_ku_edge_index,
            gcn_edge_weight_ku=data.graph_ku_edge_weight,
            temperature=0.05,
            mask_min=0.1,
            mask_max=0.25
        )

        print(f"   ✅ GMAE创建成功")
        print(f"   - 总参数数量: {sum(p.numel() for p in gmae.parameters())}")

        # 测试前向传播
        x_u_feat = data.user_interest_features
        x_e_emb = torch.randn(300, 128)

        loss, z_u, z_e_all = gmae(x_u_feat, x_e_emb)

        print(f"   ✅ GMAE前向传播成功")
        print(f"   - 损失值: {loss.item():.6f}")
        print(f"   - 用户嵌入形状: {z_u.shape}")
        print(f"   - 实体嵌入形状: {z_e_all.shape}")

        # 检查损失是否为0.6931
        if abs(loss.item() - 0.6931) < 0.001:
            print("   ⚠️  警告: 损失接近0.6931 (随机预测值)")
        else:
            print(f"   ✅ 损失正常: {loss.item():.6f}")

        # 4. 配置检查
        print("\n4. 配置检查:")

        # 检查是否使用了正确的解码器
        decoder_type = type(gmae.decoder).__name__
        print(f"   - 使用的解码器: {decoder_type}")

        if hasattr(gmae.decoder, 'decoder'):
            print(f"   ✅ 使用的是增强版解码器")
        else:
            print(f"   ❌ 可能使用的是旧版解码器")

        # 检查掩码节点数量
        print(f"   - 掩码节点数量: {masked_indices.numel()}")
        print(f"   - 目标特征形状: {target_masked.shape}")
        print(f"   - 重建特征形状: {recon_masked.shape}")

        print("\n=== 诊断完成 ===")
        return True

    except Exception as e:
        print(f"❌ 诊断失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def check_training_config():
    """检查训练配置"""
    print("\n=== 训练配置检查 ===")

    try:
        from src.parser import parse_args

        # 创建完整的测试配置
        test_config = """
data:
  dataset: "book-crossing"
  path: "data/"
  train_file: "train.txt"
  test_file: "test.txt"
  kg_file: "kg_final.txt"
  item_meta_file: "item_text.json"
  n_users: 6616
  n_items: 8853

preprocessing:
  user_enhance:
    llm_prompt_max_history: 30
    output_file: "user_interest_clustered.txt"
    cluster_num: 350
    cluster_type: 'st'
    min_cluster_size: 2
    max_cluster_ratio: 0.2
  item_enhance:
    output_embed_file: "item_sbert_embeddings.pt"
    output_graph_file: "item_semantic_graph.pt"
    sbert_model_name: 'princeton-nlp/sup-simcse-roberta-large'
    semantic_graph_k: 30

training:
  model_name: "SDKR-Test"
  embedding_dim: 128
  n_layers: 3
  epochs: 100
  batch_size: 512
  learning_rate: 0.001
  alpha: 0.5
  beta: 0.3
  gamma: 0.2
  bpr_loss_weight: 1.0
  recon_loss_weight: 0.3
  use_attention_fusion: true
"""

        # 写入临时文件
        with open("test_config_diagnose.yaml", "w") as f:
            f.write(test_config)

        # 模拟命令行参数
        import sys
        sys.argv = ['diagnose.py', '--config_file', 'test_config_diagnose.yaml',
                    '--dataset', 'book-crossing', '--gpu_id', '0', '--seed', '2024']

        config = parse_args()

        print("✅ 配置解析成功")
        print(f"   - alpha: {getattr(config, 'alpha', '未找到')}")
        print(f"   - beta: {getattr(config, 'beta', '未找到')}")
        print(f"   - gamma: {getattr(config, 'gamma', '未找到')}")
        print(f"   - recon_loss_weight: {getattr(config, 'recon_loss_weight', '未找到')}")

        # 清理
        os.remove("test_config_diagnose.yaml")

        return True

    except Exception as e:
        print(f"❌ 配置检查失败: {e}")
        return False

def main():
    """主诊断函数"""
    print("🚀 开始GMAE修复诊断")
    print("=" * 50)

    # 检查PyTorch版本
    print(f"PyTorch版本: {torch.__version__}")
    print(f"CUDA可用: {torch.cuda.is_available()}")

    tests = [
        diagnose_gmae,
        check_training_config
    ]

    results = []
    for test_func in tests:
        try:
            result = test_func()
            results.append(result)
        except Exception as e:
            print(f"❌ {test_func.__name__} 测试异常: {e}")
            results.append(False)

    print("\n" + "=" * 50)
    print("📊 诊断结果汇总:")

    passed = sum(results)
    total = len(results)

    for i, (test_func, result) in enumerate(zip(tests, results)):
        status = "✅ 通过" if result else "❌ 失败"
        print(f"  {i+1}. {test_func.__name__}: {status}")

    print(f"\n🎯 总体结果: {passed}/{total} 测试通过")

    if passed == total:
        print("🎉 GMAE修复验证通过！")
        print("\n💡 运行建议:")
        print("  1. 使用激进修复配置: --config radical-fixed")
        print("  2. 监控GMAE重建损失是否开始下降")
        print("  3. 如果损失仍为0.6931，说明有根本性问题")
    else:
        print("⚠️  诊断发现问题，请检查相关修复。")

    return passed == total

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)